﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using OpenCLNet;

namespace GPUTracer
{
    public static class KDTreeBuilder
    {
        private const int MaxPrimitives = 3;
        private const float EmptyBonus = 0.2f;
        private const float isectCost = 80;
        private const float traversalCost = 1;

        private struct BoundingEdge
        {
            public float Pos;
            public int Index;
            public bool Start;
        }

        public static KDTreeElement Build(Object[] objects)
        {
            var ob = new List<Object>();
            foreach (var o in objects)
            {
                if (o.GetBoundingBox().HasValue)
                    ob.Add(o);
            }

            BoundingBox objectsBBox = GetObjectsBoundingBox(ob);
            KDTreeElement root = SplitObjects(ob.ToArray(), objectsBBox);

            return root;
        }

        private static KDTreeElement SplitObjects(Object[] objects, BoundingBox bbox)
        {
            if (objects.Count() <= MaxPrimitives)
            {
                return new KDTreeLeaf(objects, bbox);
            }

            float bestCost = float.MaxValue;
            int bestIndex = 0;
            float bestPos = 0.0f;

            //Choose axis to split along:
            Axis axis;
            var objectsSize = new Float4(bbox.Max.S0 - bbox.Min.S0,
                                         bbox.Max.S1 - bbox.Min.S1,
                                         bbox.Max.S2 - bbox.Min.S2, 0);
            float totalSA = 2 * (objectsSize.S0 * objectsSize.S1 + objectsSize.S0 * objectsSize.S2 + objectsSize.S1 * objectsSize.S2);
            float invTotalSA = 1/totalSA;

            if (objectsSize.S0 > objectsSize.S1 && objectsSize.S0 > objectsSize.S2)
                axis = Axis.XAxis;
            else if (objectsSize.S1 > objectsSize.S2)
                axis = Axis.YAxis;
            else
                axis = Axis.ZAxis;

            int retries = 2;

            while (retries > 0)
            {
                List<BoundingEdge> edges;
                bestCost = SearchSplittingPosition(bbox, objectsSize, axis, objects, invTotalSA, bestCost, ref bestPos,
                                                   ref bestIndex, out edges);

                if (bestCost < float.MaxValue)
                {
                    KDTreeNode node = BuildNode(bbox, axis, objects, bestPos, bestIndex, edges);
                    return node;
                }
                
                if (retries > 0)        //try next axis
                {
                    axis = (Axis) (((int)(axis + 1))%3);
                    retries--;
                }
            }

            return new KDTreeLeaf(objects, bbox);
        }

        private static KDTreeNode BuildNode(BoundingBox bbox, Axis axis, Object[] objects, float bestPos, int bestIndex, List<BoundingEdge> edges)
        {
            var node1Objects = new List<Object>();
            var node2Objects = new List<Object>();

            var node1BB = new BoundingBox(bbox.Min,
                                          new Float4(axis == Axis.XAxis ? bestPos : bbox.Max.S0,
                                                     axis == Axis.YAxis ? bestPos : bbox.Max.S1,
                                                     axis == Axis.ZAxis ? bestPos : bbox.Max.S2, 0));
            var node2BB = new BoundingBox(new Float4(axis == Axis.XAxis ? bestPos : bbox.Min.S0,
                                                     axis == Axis.YAxis ? bestPos : bbox.Min.S1,
                                                     axis == Axis.ZAxis ? bestPos : bbox.Min.S2, 0),
                                          bbox.Max);

            for (int i = 0; i < bestIndex; i++)
            {
                if (edges[i].Start)
                    node1Objects.Add(objects[edges[i].Index]);
            }
            for (int i = bestIndex + 1; i < edges.Count(); i++)
            {
                if (!edges[i].Start)
                    node2Objects.Add(objects[edges[i].Index]);
            }

            var node = new KDTreeNode(axis, bestPos);
            node.BBox = bbox;
            node.LeftChild = SplitObjects(node1Objects.ToArray(), node1BB);
            node.RightChild = SplitObjects(node2Objects.ToArray(), node2BB);
            node.NumObjects = objects.Count();
            return node;
        }

        private static float SearchSplittingPosition(BoundingBox bbox, Float4 objectsSize, Axis axis, Object[] objects, float invTotalSA, float bestCost, ref float bestPos, ref int bestIndex, out List<BoundingEdge> edges)
        {
            var otherAxis1 = (axis == Axis.XAxis) ? Axis.YAxis : Axis.XAxis;
            var otherAxis2 = (axis == Axis.YAxis) ? Axis.ZAxis : Axis.YAxis;

            //Search for a potential splitting plane along axis:
            edges = new List<BoundingEdge>();
            int c = 0;
            foreach (var o in objects)
            {
                var bb = o.GetBoundingBox();
                edges.Add(new BoundingEdge() { Pos = GetVectorAxisElement(bb.Value.Min, axis), Start = true, Index = c});
                edges.Add(new BoundingEdge() { Pos = GetVectorAxisElement(bb.Value.Max, axis), Start = false, Index = c });
                c++;
            }
            edges.Sort((x,y) => x.Pos.CompareTo(y.Pos));

            int nBelow = 0;
            int nAbove = objects.Count();
            c = 0;
            foreach (var boundingEdge in edges)
            {
                if (!boundingEdge.Start)
                    nAbove--;

                float edget = boundingEdge.Pos;
                if (edget > GetVectorAxisElement(bbox.Min, axis) && edget < GetVectorAxisElement(bbox.Max, axis))
                {
                    //compute cost of splitting here:
                    float belowSA = 2 *
                                    (GetVectorAxisElement(objectsSize, otherAxis1) * GetVectorAxisElement(objectsSize, otherAxis2) +
                                     (edget - GetVectorAxisElement(bbox.Min, axis)) *
                                     (GetVectorAxisElement(objectsSize, otherAxis1) + GetVectorAxisElement(objectsSize, otherAxis2)));
                    float aboveSA = 2 *
                                    (GetVectorAxisElement(objectsSize, otherAxis1) * GetVectorAxisElement(objectsSize, otherAxis2) +
                                     (GetVectorAxisElement(bbox.Max, axis) - edget) *
                                     (GetVectorAxisElement(objectsSize, otherAxis1) + GetVectorAxisElement(objectsSize, otherAxis2)));
                    float pBelow = belowSA * invTotalSA;
                    float pAbove = aboveSA * invTotalSA;
                    float eb = (nAbove == 0 || nBelow == 0) ? EmptyBonus : 0;
                    float cost = traversalCost + isectCost * (1 - eb) * (pBelow * nBelow + pAbove * nAbove);

                    if (cost < bestCost)
                    {
                        bestCost = cost;
                        bestIndex = c;
                        bestPos = boundingEdge.Pos;
                    }
                }

                if (boundingEdge.Start)
                    nBelow++;
                c++;
            }
            return bestCost;
        }

        private static BoundingBox GetObjectsBoundingBox(IEnumerable<Object> objects)
        {
            var mS0 = float.MaxValue;
            var mS1 = float.MaxValue;
            var mS2 = float.MaxValue;
            var nS0 = float.MinValue;
            var nS1 = float.MinValue;
            var nS2 = float.MinValue;

            foreach (var o in objects)
            {
                var obb = o.GetBoundingBox();
                if (obb.HasValue)
                {
                    if (obb.Value.Min.S0 < mS0)
                        mS0 = obb.Value.Min.S0;
                    if (obb.Value.Min.S1 < mS1)
                        mS1 = obb.Value.Min.S1;
                    if (obb.Value.Min.S2 < mS2)
                        mS2 = obb.Value.Min.S2;

                    if (obb.Value.Max.S0 > nS0)
                        nS0 = obb.Value.Max.S0;
                    if (obb.Value.Max.S1 > nS1)
                        nS1 = obb.Value.Max.S1;
                    if (obb.Value.Max.S2 > nS2)
                        nS2 = obb.Value.Max.S2;
                }
            }

            return new BoundingBox(new Float4(mS0, mS1, mS2, 0), new Float4(nS0, nS1, nS2, 0));
        }

        private static float GetVectorAxisElement(Float4 f, Axis axis)
        {
            switch (axis)
            {
                case Axis.XAxis:
                    return f.S0;
                case Axis.YAxis:
                    return f.S1;
                case Axis.ZAxis:
                    return f.S2;
                default:
                    throw new ArgumentOutOfRangeException("axis");
            }
        }
    }
}
